from typing import Iterable

from pyflink.common import SimpleStringSchema, WatermarkStrategy, Time, Types
from pyflink.datastream import StreamExecutionEnvironment, ProcessFunction, KeyedProcessFunction, ProcessWindowFunction
from pyflink.datastream.connectors.kafka import KafkaSource, KafkaOffsetsInitializer
from pyflink.datastream.functions import KEY, IN, OUT
from pyflink.datastream.window import TumblingProcessingTimeWindows

# 1、创建flink执行环境
env = StreamExecutionEnvironment.get_execution_environment()

env.set_parallelism(1)

source = KafkaSource.builder() \
    .set_bootstrap_servers("master:9092") \
    .set_topics("bigdata") \
    .set_group_id("my-group") \
    .set_starting_offsets(KafkaOffsetsInitializer.latest()) \
    .set_value_only_deserializer(SimpleStringSchema()) \
    .build()

kafka_ds = env.from_source(source, WatermarkStrategy.no_watermarks(), "Kafka Source")

class FlatMapProcessFunction(ProcessFunction):
    def process_element(self, line, ctx: 'ProcessFunction.Context'):
        processing_time = ctx.timer_service().current_processing_time()
        print(processing_time)

        words = line.split(',')
        for word in words:
            yield word


words_ds = kafka_ds.process(FlatMapProcessFunction())

key_by_ds = words_ds.key_by(lambda kv: kv)

class ReduceKeyedProcessFunction(KeyedProcessFunction):

    def __init__(self):
        self.word_count = dict()
    def process_element(self, word, ctx: 'KeyedProcessFunction.Context'):
        if word in self.word_count:
            self.word_count[word] += 1
        else:
            self.word_count[word] = 1

        yield word, self.word_count[word]


count_ds = key_by_ds.process(ReduceKeyedProcessFunction())

window_ds = key_by_ds.window(TumblingProcessingTimeWindows.of(Time.seconds(5)))

class WindowProcessWindowFunction(ProcessWindowFunction):
    def process(self, key: KEY, context: 'ProcessWindowFunction.Context', elements: Iterable[IN]) -> Iterable[OUT]:
        print(list(elements))

        count = 0
        for element in elements:
            count += 1

        yield key, count


window_count_ds = window_ds.process(WindowProcessWindowFunction())


window_count_ds.print()

env.execute()